#include "DepthTracker.h"

using namespace cv;

DepthTracker::DepthTracker(void)
{
	tmpBins = NULL;
	subMat = NULL;
	depthRange.x = 0;
	depthRange.y = 32767;
	inRangeMat = NULL;
	isTracking = false;
}

DepthTracker::~DepthTracker(void)
{
	delete[] tmpBins;
	if (subMat != NULL) cvReleaseMatHeader(&subMat);
	if (inRangeMat != NULL) cvReleaseMat(&inRangeMat);
}

void DepthTracker::showHist(CvMat *mat, CvRect subRect) {
	CvMat *pGrayImg16U = cvCreateMat(mat->rows, mat->cols, CV_16UC1);
	cvConvert(mat, pGrayImg16U);

	subRect.width = rect.width;
	subRect.height = rect.height;
	cvGetSubRect(pGrayImg16U, subMat, subRect);
	hist.calcHistFast(subMat, (short)0);
	hist.threshHist(rect.height*rect.width/100);
	cout << "debug dis: " << calcDis(hist, hist) << endl;
	hist.show("debug hist");
	cvReleaseMat(&pGrayImg16U);
}

void DepthTracker::calcRg2(CvMat *subMat) {
	hist.calcHistFast(subMat, short(0));
	center.z = hist.mainMean(0.5);

	hist.sortBins();
	hist.show("debug hist", -1, 50);
	hist.accumulate();

	depthRange.x = numeric_limits<float>::max();
	depthRange.y = numeric_limits<float>::min();
	for (int i=0; i<hist.binNum; ++i)
		if (hist.bins[i] <= 0.3 * hist.bins[hist.binNum-1]) {
			if (depthRange.x > hist.eachBinRg[i])
				depthRange.x = hist.eachBinRg[i];
			if (depthRange.y < hist.eachBinRg[i])
				depthRange.y = hist.eachBinRg[i];
		}
		else {
			break;
		}
}

void DepthTracker::calcRg(CvMat *subMat) {
	short *sPtr;

	depthRange.x = numeric_limits<float>::max();
	depthRange.y = numeric_limits<float>::min();
	bool startRg=false;
	const int TH = 500;

	CvPoint2D32f localRg = depthRange;

	for (int i=0; i<subMat->rows; ++i) {
		sPtr = (short*)cvPtr2D(subMat, i, 0);

		int lastVal = -1;

		for (int j=0; j<subMat->width; ++j) {
			if (sPtr[j] == 0) continue;
			if (lastVal == -1) {
				localRg.x = localRg.y = sPtr[j];
				startRg = true;
			}
			else {
				if (lastVal-sPtr[j] > TH) {
					localRg.x = localRg.y = sPtr[j];
					startRg = true;
				}
				else if (startRg) {
					if (sPtr[j]-lastVal > TH) {
						startRg = false;
						if (depthRange.x>localRg.x) depthRange.x = localRg.x;
						if (depthRange.y<localRg.y) depthRange.y = localRg.y;
					}
					else {
						if (localRg.x>sPtr[j]) localRg.x = sPtr[j];
						if (localRg.y<sPtr[j]) localRg.y = sPtr[j];
					}
				}
			}
			lastVal = sPtr[j];
		}
		if (startRg) {
			startRg = false;
			if (depthRange.x>localRg.x) depthRange.x = localRg.x;
			if (depthRange.y<localRg.y) depthRange.y = localRg.y;
		}
	}

	cout << depthRange.x << " " << depthRange.y << endl;


}

void DepthTracker::initTracking(Mat mat, Rect rect) {
	inRangeMat = cvCreateMat(mat.rows, mat.cols, CV_16SC1);
	subMat = cvCreateMatHeader(rect.height, rect.width, mat.type());

	this->rect = rect;
	int ub2 = 8000, lb2 = 4000, rgPerBin = 2;
	hist.setBinNum(cvCeil((ub2-lb2)/rgPerBin));
	hist.setRange(lb2, ub2);

	Point first, second;
	first = rect.tl();
	second = rect.br();

	center.x = rect.x+rect.width/2;
	center.y = rect.y+rect.height/2;

	Mat tmpSubMat = Mat(mat, rect);
	CvMat cvsubMat = tmpSubMat;
	calcRg(&cvsubMat);

	cout << " depth range: " << depthRange.x << " " << depthRange.y << endl;
	lb = depthRange.x - upperOffset;
	ub = depthRange.y + lowerOffset;

	isTracking = true;
}

void DepthTracker::initTracking(CvMat *mat, CvRect rect, double *depthVal) {
	this->rect = rect;
	inRangeMat = cvCreateMat(mat->rows, mat->cols, CV_16SC1);
	int ub2 = 8000, lb2 = 4000, rgPerBin = 2;
	hist.setBinNum(cvCeil((ub2-lb2)/rgPerBin));
	hist.setRange(lb2, ub2);

	center.x = rect.x+rect.width/2;
	center.y = rect.y+rect.height/2;

	center.z = cvGetReal2D(mat, center.y, center.x);

	subMat = cvCreateMatHeader(rect.height, rect.width, mat->type);
	cvGetSubRect(mat, subMat, rect);

	calcRg(subMat);

	cout << "center.z: " << center.z << " depth range: " << depthRange.x << " " << depthRange.y << endl;
	lb = depthRange.x - upperOffset;
	ub = depthRange.y + lowerOffset;

// 	applyDepthRange(mat);
// 	cvShowImage("rg", inRangeMat);
// 	cvGetSubRect(inRangeMat, subMat, rect);
// 	cvShowImage("rgSub", subMat);
// 	hist.calcHistFast(subMat, short(0));
// 	hist.sortBins();
// 	hist.show("debug hist", -1, 50);
// 	cvWaitKey(0);
// 	center.z = hist.mainMean(1);


	if (depthVal != NULL)
		*depthVal = center.z;

	isTracking = true;
}

double DepthTracker::calcDis(const SimpleHistogram &a, const SimpleHistogram& b) {
	double mean=0, varDev=0;
	int binNum = a.binNum;

	int cnt = 0;
	for (int i=0; i<binNum; ++i)
		if (a.bins[i] != 0) {
			mean += a.bins[i]*(i+0.5);
			cnt += a.bins[i];
		}
		mean = mean / cnt;

		double mean2 = 0;
		cnt = 0;
		for (int i=0; i<binNum; ++i)
			if (b.bins[i] != 0) {
				mean2 += b.bins[i]*(i+0.5);
				cnt += b.bins[i];
			}
			mean2 = mean2 / cnt;

			cout << "mean: " << mean << " mean2: " << mean2 << endl;
			double binW = (a.ranges[1]-a.ranges[0])/a.binNum;
			int offset = cvRound((mean - mean2));
			cout << "offset: " << offset << endl;
			mean = 0;
			cnt = 0;
			for (int i=0; i<binNum; ++i) {
				if (i+offset<binNum && 0<=i+offset) {
					tmpBins[i] = a.bins[i] - b.bins[i+offset];
					if (a.bins[i] != 0 || b.bins[i+offset] != 0) {
						mean += tmpBins[i];
						++cnt;
					}
				}
				else
					tmpBins[i] = 0;
			}

			cout << "cnt: " << cnt << " mean: " << mean << endl;
			if (cnt == 0) mean = 0;
			else mean /= cnt;
			varDev = 0;
			for (int i=0; i<binNum; ++i) {
				varDev += (tmpBins[i]-mean)*(tmpBins[i]-mean);
			}
			return varDev;
}


typedef struct TRectInfo {
public:
	Rect rect;
	Mat canvas;
	string name;
} TRectInfo;

void drawRect(int type, int x, int y, int flag, void* param) {
	TRectInfo *rectInfo = (TRectInfo*)param;
	Point first, second;
	first = rectInfo->rect.tl();
	second = rectInfo->rect.br();

	switch (type) {
	case CV_EVENT_LBUTTONDOWN:
		first.x = x;
		first.y = y;
		second = first;
		break;
	case CV_EVENT_LBUTTONUP:
		{
			second.x = x;
			second.y = first.y+x-first.x;
			Mat canvas;
			rectInfo->canvas.copyTo(canvas);
			char words[100];
			sprintf_s(words, 100, "first=(%d, %d), second=(%d, %d)", first.x, first.y, second.x, second.y);
			int maxInt = (1 << canvas.type()*8) - 1;
			putText(canvas, words, Point(50,50), CV_FONT_HERSHEY_SIMPLEX,
				0.7, Scalar(maxInt));
			rectangle(canvas, first, second, Scalar(maxInt));
			imshow(rectInfo->name, canvas);
		}
		break;
	case CV_EVENT_MOUSEMOVE:
		if (flag & CV_EVENT_FLAG_LBUTTON) {
			second.x = x;
			second.y = first.y+x-first.x;
			Mat canvas;
			rectInfo->canvas.copyTo(canvas);
			char words[100];
			sprintf_s(words, 100, "first=(%d, %d), second=(%d, %d)", first.x, first.y, second.x, second.y);
			int maxInt = (1 << canvas.type()*8) - 1;
			putText(canvas, words, Point(50,50), CV_FONT_HERSHEY_SIMPLEX,
				0.7, Scalar(maxInt));
			rectangle(canvas, first, second, Scalar(maxInt));
			imshow(rectInfo->name, canvas);
		}
		break;
	default:
		break;
	}
	rectInfo->rect = Rect(first, second);
}


Rect DepthTracker::tracking(const Mat &mat) {
	if (isTracking) {
		CvMat cvmat = mat;
		return tracking(&cvmat);
	}
	else {
		TRectInfo depthRect;
		depthRect.name = "init_tracker_depth";

		namedWindow("init_tracker_depth");
		setMouseCallback("init_tracker_depth", drawRect, &depthRect);

		char c=0;
		while (true) {
			imshow("init_tracker_depth", mat);
			depthRect.canvas = mat;
			c = cvWaitKey(0);
			if (c == 13) break;
			if (c == 27) break;
		}
		destroyWindow("init_tracker_depth");
		if (c == 27) return depthRect.rect;
		
		initTracking(mat, depthRect.rect);
		return depthRect.rect;
	}
}

CvRect DepthTracker::tracking(CvMat *mat, double *depthVal) {
	CvRect subRect;
	applyDepthRange(mat);
	cvShowImage("ranged", inRangeMat);

	subRect.x = max(0, int(center.x-rect.width/2-offset));
	subRect.width = min(mat->cols-subRect.x, rect.width+2*offset);
	subRect.y = max(0, int(center.y-rect.height/2-offset));
	subRect.height = min(mat->rows-subRect.y, rect.height+2*offset);

	cvGetSubRect(inRangeMat, subMat, subRect);


	
	
	int len = cvRound(depthRange.y-depthRange.x);

	SimpleHistogram hist;
	hist.setBinNum(len);
	hist.setRange(cvRound(depthRange.x), cvRound(depthRange.y));
	hist.calcHistFast(subMat, short(0));

// 	hist.show("debug", 1);
// 	namedWindow("backmap", 0);
// 	imshow("backmap", hist.backmapColor(subMat, short(0)));
// 
// 	cvWaitKey(0);
	hist.sortBins();



	short *sPtr;
	int cnt = 0;
	center.x = 0;
	center.y = 0;
	vector<Point2f> samples;
	double largeDis = norm(Point2f(subMat->rows/2, subMat->cols/2));
	double sum = 0;

	for (int i=0; i<subMat->rows; ++i) {
		sPtr = (short*) cvPtr2D(subMat, i, 0);
		for (int j=0; j<subMat->cols; ++j, ++sPtr) {
			short val = *sPtr;
			if (val == 0) continue;
			double low = hist.eachBinRg[0]+hist.ranges[0]-1000;
			double high = hist.eachBinRg[0]+hist.binW+hist.ranges[0];
			if (low <= val && val < high) {
				double coef = (1 - 0.5*norm(Point2f(i-subMat->rows/2,j-subMat->cols/2))/largeDis);
// 				center.x += j;
// 				center.y += i;
				++cnt;
				sum += coef;
				center.x += j*coef;
				center.y += i*coef;
				samples.push_back(Point2f(j,i));
			}
		}
	}

	cout << "center: " << center.x/sum << " " << center.y/sum << endl;
	
// 	center.x = subRect.x + center.x/cnt;
// 	center.y = subRect.y + center.y/cnt;
	center.x = subRect.x + center.x/sum;
	center.y = subRect.y + center.y/sum;

	cout << "sample size: " << samples.size() << endl;
	Mat drawSample = Mat::zeros(subMat->rows, subMat->cols, CV_8U);
	
	for (int i=0; i<samples.size(); ++i) {
		drawSample.at<uchar>(samples[i].y, samples[i].x) = 255;
	}

// 	namedWindow("sample", 0);
// 	imshow("sample", drawSample);



	center.z = cvGetReal2D(mat, center.y, center.x);

	calcRg(subMat);

// 	cvWaitKey(0);
// 	cout << "center.z: " << center.z << " depthrange: " << depthRange.x << " " << depthRange.y << endl;
	lb = depthRange.x - lowerOffset;
	ub = depthRange.y + upperOffset;

// 	applyDepthRange(mat);
// 	cvGetSubRect(inRangeMat, subMat, rect);
// 	hist.calcHistFast(subMat, short(0));
// 	center.z = hist.mainMean(1);

	if (depthVal != NULL)
		*depthVal = center.z;
	subRect.x = max(0, int(center.x-rect.width/2));
	subRect.width = min(rect.width, int(center.x+rect.width/2));
	subRect.y = max(0, int(center.y-rect.height/2));
	subRect.height = min(rect.height, int(center.y+rect.height/2));

	return subRect;
}

void DepthTracker::applyDepthRange(CvMat *mat)
{
	int cnt = 0;
	short* matPtr;
	short* resPtr;
	matPtr = mat->data.s;
	resPtr = inRangeMat->data.s;
	for (int i=0; i<mat->rows; ++i)
		for (int j=0; j<mat->cols; ++j, ++matPtr, ++resPtr) {
			if (*matPtr<lb || *matPtr>ub) *resPtr = 0;
			else {
				*resPtr = *matPtr;
				++cnt;
			}
		}

	if (cnt < atLeastPixel) {
		lb -= lowerOffset;
		ub += upperOffset;
		applyDepthRange(mat);
	}
}

